
from llama_index.core.indices.property_graph import SimpleLLMPathExtractor
from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm

from llama_index.core import SQLDatabase
from llama_index.core.indices.struct_store import SQLStructStoreIndex
from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer

# 1. 创建SQLite内存数据库
engine = create_engine("sqlite:///:memory:")
metadata_obj = MetaData()

# 2. 定义表结构
table_name = "city_stats"
city_stats_table = Table(
    table_name,
    metadata_obj,
    Column("city_name", String(16), primary_key=True),
    Column("population", Integer),
    Column("country", String(16), nullable=False)
)
metadata_obj.create_all(engine)

# 3. 插入示例数据
data = [
    {"city_name": "Toronto", "population": 2731571, "country": "Canada"},
    {"city_name": "Tokyo", "population": 13929286, "country": "Japan"},
    {"city_name": "Berlin", "population": 3645000, "country": "Germany"}
]
with engine.connect() as conn:
    for row in data:
        conn.execute(city_stats_table.insert().values(**row))
    conn.commit()

# 4. 创建SQLDatabase包装器
sql_database = SQLDatabase(engine)

# 5. 构建SQLStructStoreIndex
index = SQLStructStoreIndex.from_documents(
    documents=[],  # 空文档列表，因为数据来自SQL
    sql_database=sql_database,
    table_name=table_name,
)

# 6. 创建查询引擎
query_engine = index.as_query_engine()
response = query_engine.query("哪个城市人口超过1000万？")
print(response)
